import numpy as np
import scipy.io as scio
import torch
from torch.utils.data import DataLoader
from data_class import *
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
import torch
from torchvision.models import *
import matplotlib.pyplot as plt
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image
# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

# 读取mat, 创建数据集
crack_db = scio.loadmat("crack_db.mat")['crack_db']
crack_st = scio.loadmat("crack_test.mat")['crack_test']
image_train = crack_db[0,0]['S']      # 数据库的散射矩阵
target_train = crack_db[0,0]['size']  # 数据库的裂纹长度
image_test = crack_st[0,0]['S']
target_test = crack_st[0,0]['size']
train_data = (image_train, target_train)    # 创建元组
test_data = (image_test, target_test)

train_dataset = MyDataset(train_data,720)       # 创建训练集
test_dataset = MyDataset(test_data, 684)        # 创建测试集
train_dataloader = DataLoader(train_dataset,batch_size=10)
test_dataloader = DataLoader(test_dataset,batch_size=10)

# 神经网络模型
net = torchvision.models.resnet18(resnet.BasicBlock,pretrained=True)
net.conv1 = nn.Conv2d(in_channels=1,out_channels=64,kernel_size=7,stride=1,padding=3)

set_parameter_require_grad(net,True)    # 冻结除全连接以外的层

net = net.to(device)
# 损失函数
loss_fn = nn.MSELoss()
loss_fn = loss_fn.to(device)
# 创建优化器
learn_rate = 0.1
optimizer = torch.optim.SGD(net.parameters(),lr=learn_rate)
# 设置训练网络的参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练轮数
epoch = 3
batch_size = 10
# 使用tensorboard观测
writer = SummaryWriter("./logs_train")
# 测试模型并记录损失值
losses = []
max_loss_index = 0
max_loss_value = float('-inf')
for i in range(epoch):
    print("第{}轮训练开始".format(i+1))
    # 训练步骤
    for data in train_dataloader:
        imgs, targets = data        # img是散射矩阵，target是裂纹长度
        imgs = imgs.to(device)
        targets = targets.to(device)
        output = net(imgs)
        targets = targets.to(torch.float32)
        loss = loss_fn(output,targets)  # 训练的损失
        # 优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step += 1
        if total_train_step % 50 == 0:
            print("训练次数:{},loss:{}".format(total_train_step,loss))
            # 用折线图观测
            writer.add_scalar("tran_loss",loss.item(),total_train_step)
    # 测试步骤
    total_test_loss = 0
    with torch.no_grad():
        for j,data in enumerate(test_dataloader):
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            outputs = net(imgs)
            loss = loss_fn(outputs,targets)
            if i == epoch -1:
                max_loss_index,max_loss_value = find_max_loss(j,losses,loss,max_loss_value,max_loss_index)
            total_test_loss = total_test_loss + loss*batch_size
    total_test_loss = total_test_loss/684
    print("整体测试集上的loss:{}".format(total_test_loss.item()))  # 测试的损失
    writer.add_scalar("test_loss",total_test_loss,total_test_step)
    total_test_step += 1
# print(train_dataset[0][0])    # 第1个散射矩阵，第二个维度是散射矩阵/裂长度
# 输入是散射矩阵，与数据库比对，得到裂纹的长度和取向角
# target标签是720个裂纹的size
# 720个tensor，每个tensor大小为(1,91,91)
# 打印最大损失值及其索引
print("Maximum loss value:", max_loss_value)
print("Index of maximum loss:", max_loss_index)
# 获取对应的数据
max_loss_data = test_dataset[max_loss_index][0]
matrix = torch.squeeze(max_loss_data[0],dim=0)      # 删去一个维度
plt.imshow(matrix, cmap='viridis')  # 使用 'viridis' 色彩映射
plt.colorbar()  # 添加颜色条
plt.show()
writer.close()